use numpy::ndarray::Array1;
use numpy::{IntoPyArray, PyArrayDyn};
use pyo3::{pymodule, types::PyModule, PyResult, Python, Bound};

use rand::prelude::*;
use rand::Rng;
use statrs::distribution::{Cauchy, Continuous};
use std::f64::consts::PI;

#[pymodule]
fn py_sinc_rand<'py>(_py: Python<'py>, m: &Bound<'py, PyModule>) -> PyResult<()> {
    fn sinc(x: f64) -> f64 {
        if x == 0.0 { 1.0 } else { x.sin() / x }
    }
    
    fn sinc_squared_pdf(x: f64, location: f64, scale: f64) -> f64 {
        let psf = (1.0 / (scale * PI).sqrt()) * sinc((x - location) / scale);
        psf * psf
    }
    
    fn cauchy_pdf(x: f64, location: f64, scale: f64) -> f64 {
        let dist = Cauchy::new(location, scale).unwrap();
        dist.pdf(x)
    }


    #[pyfn(m)]
    #[pyo3(name = "sinc_rand")]
    fn sinc_rand_py<'py>(py: Python<'py>, location: f64, scale: f64 , size: usize) -> Bound<'py, PyArrayDyn<f64>> {
        let sigma = scale;
        let mu = location;
        let m = 3.0;
        let proposal_sigma = 2.0 * sigma;
    
        let mut rng = rand::thread_rng();
        let mut samples = Vec::with_capacity(size);
    
        let cauchy = Cauchy::new(mu, proposal_sigma).unwrap();
    
        while samples.len() < size {
            let to_generate: usize = size - samples.len();
    
            let xs: Vec<f64> = (0..to_generate).map(|_| cauchy.sample(&mut rng)).collect();
    
            for &x in &xs {
                let u: f64 = rng.r#gen();
                let target: f64 = sinc_squared_pdf(x, mu, sigma);
                let proposal: f64 = cauchy_pdf(x, mu, proposal_sigma);
                if u < target / (m * proposal) {
                    samples.push(x);
                }
            }
        }
        let array = Array1::from(samples).into_dyn();
        array.into_pyarray(py)
    }

    Ok(())
}